from arcpy.sa import *
import arcpy
from stormwaterutilities import addMessage
import stormwaterutilities
import accumulatetonetwork
import tempfile
import os
import numpy

import pollutantdata as pd

import rawq_susceptibility
import rawq_decumulate

def CreateWQNetwork(outputGdb, linkName, nodeName, catchments, flowDir, accum, weight):

    drainageDir = r'\\ws18325\Stormwater D\RiskModel\Drainage'
    tempDir = r"d:\temp"

    arcpy.env.overwriteOutput = True

    # Copy drainage network
    sourcePath = os.path.join(drainageDir, 'DrainageNetwork.gdb')

    if arcpy.Exists(outputGdb):
        addMessage('Deleting existing drainage network database for latest RAWQ...')
        gdb = arcpy.Delete_management(outputGdb)
        addMessage('Deleted: {0}'.format(gdb))

    addMessage('Copying drainage network database from {0}'.format(sourcePath))
    errors = arcpy.Copy_management(sourcePath, outputGdb)
    addMessage('Copied to: {0}'.format(errors))

    arcpy.Rename_management(os.path.join(outputGdb, 'StormDrainConveyance'), linkName)
    arcpy.Rename_management(os.path.join(outputGdb, 'StormDrainNodes'), nodeName)

    # Copy flow direction grid
    sourceFlowDir = os.path.join(drainageDir, 'flowdir_tau.tif')
    destFlowDir = os.path.join(outputDir, 'flowdir_tau.tif')

    if arcpy.Exists(destFlowDir):
        addMessage('Deleting existing flow direction grid for latest RAWQ...')
        gdb = arcpy.Delete_management(destFlowDir)
        addMessage('Deleted: {0}'.format(gdb))

    addMessage('Copying flow direction grid from {0}'.format(sourceFlowDir))
    errors = arcpy.Copy_management(sourceFlowDir, destFlowDir)
    addMessage('Copied to: {0}'.format(errors))

    arcpy.CheckOutExtension('Spatial')

    #Read each landuse type in the landuse raster
    valueFieldName = 'Value'
    descriptiveFieldName = 'Code'
    fields =[valueFieldName, descriptiveFieldName]

    values, descriptions = zip(*[(row[0], row[1]) for row in arcpy.da.SearchCursor(accum, (fields))])
    outputs = [os.path.join(tempDir, description + '_accum.tif') for description in descriptions]

    #Accumulate each landuse type in the landuse raster
    for value, description, output in zip(values, descriptions, outputs):
        addMessage('%s %s %s'%(value, description, output))
        raster = arcpy.sa.Con(IsNull(Raster(weight)) | IsNull(Raster(accum)), 0, Con(Raster(accum) == value, Raster(weight),0))

        tempRaster = os.path.join(tempDir, 'rawq_temp.tif')

        addMessage('Saving temp raster to ' + tempRaster)
        raster.save(tempRaster)

        accumulatetonetwork.AccumulateToNetwork(
            os.path.join(outputGdb, linkName),
            os.path.join(outputGdb, nodeName),
            catchments, flowDir,
            tempRaster, description + '_sq_ft', 9)

    try:
        arcpy.Delete_management(tempRaster)
    except:
        addMessage("Warning: Unable to delete last temporary landuse raster.")

def CalculateLoading(table):

    #Calculate pollutant loadings, scores, flow rate, and volume
    rainfallDepth = pd.annualRainfallDepth
    rainfallDuration = 365

    addMessage('Reading feature class to NumPy array')

    landuseFields = [(s + '_sq_ft') for s in pd.landuse]
    interestingLanduse = list(pd.landuse)
    interestingLanduse.remove('WAT')

    addMessage(interestingLanduse)
    inArray = arcpy.da.FeatureClassToNumPyArray(table, (['OID@', 'accumulated_area_sq_ft'] + landuseFields))

    addMessage('Calculating loadings and concentrations')
    volume = inArray['accumulated_area_sq_ft'] * rainfallDepth / 12
    flow = volume / (rainfallDuration * 24 * 60 * 60)

    numpy.seterr(divide='ignore')

    tssLoad = rainfallDepth * 5.202e-6 * sum([(inArray[l + '_sq_ft'] * pd.tssLuConcentration[l].mean) for l in interestingLanduse])
    tssConcentration = numpy.where(volume > 0, tssLoad / volume, 0) * 16018.4
    tssScore = stormwaterutilities.mapValues(tssConcentration, pd.tssScoreBasis, pd.scores)

    pbLoad = rainfallDepth * 5.202e-9 * sum([(inArray[l + '_sq_ft'] * pd.pbLuConcentration[l].mean) for l in interestingLanduse])
    pbConcentration = numpy.where(volume > 0, pbLoad / volume, 0) * 16018400
    pbScore = stormwaterutilities.mapValues(pbConcentration, pd.pbScoreBasis, pd.scores)

    tpLoad = rainfallDepth * 5.202e-6 * sum([(inArray[l + '_sq_ft'] * pd.tpLuConcentration[l].mean) for l in interestingLanduse])
    tpConcentration = numpy.where(volume > 0, tpLoad / volume, 0) * 16018.4
    tpScore = stormwaterutilities.mapValues(tpConcentration, pd.tpScoreBasis, pd.scores)

    bodLoad = rainfallDepth * 5.202e-6 * sum([(inArray[l + '_sq_ft'] * pd.bodLuConcentration[l].mean) for l in interestingLanduse])
    bodConcentration = numpy.where(volume > 0, bodLoad / volume, 0) * 16018.4
    bodScore = stormwaterutilities.mapValues(bodConcentration, pd.bodScoreBasis, pd.scores)

    ecoliLoad = rainfallDepth * 23.597 * sum([(inArray[l + '_sq_ft'] * pd.ecoliLuConcentration[l].mean) for l in interestingLanduse])
    ecoliConcentration = numpy.where(volume > 0, ecoliLoad / volume, 0) * 0.00353146667
    ecoliScore = stormwaterutilities.mapValues(ecoliConcentration, pd.ecoliScoreBasis, pd.scores)

    numpy.seterr(divide='warn')

    #For now, magnitude score is unweighted mean of pollutants
    magnitudeScore = (tssScore + pbScore + tpScore + bodScore + ecoliScore) / 5

    addMessage('Writing loadings and concentrations to NumPy array')
    typeList = numpy.dtype([('oid', numpy.int32), ('volume_cu_ft', numpy.float32),('flow_cfs', numpy.float32),
                            ('TSS_lbs', numpy.float32), ('TSS_mg_per_L', numpy.float32), ('TSS_score', numpy.float32),
                            ('PbD_lbs', numpy.float32), ('PbD_ug_per_L', numpy.float32), ('PbD_score', numpy.float32),
                            ('TP_lbs', numpy.float32), ('TP_mg_per_L', numpy.float32), ('TP_score', numpy.float32),
                            ('BOD_lbs', numpy.float32), ('BOD_mg_per_L', numpy.float32), ('BOD_score', numpy.float32),
                            ('ecoli_cfu', numpy.float32), ('ecoli_cfu_per_100mL', numpy.float32), ('ecoli_score', numpy.float32),
                            ('magnitude_score', numpy.float32)])
    outArray = numpy.asarray(zip(inArray['OID@'], volume, flow,
                                 tssLoad, tssConcentration, tssScore, pbLoad, pbConcentration, pbScore, tpLoad, tpConcentration, tpScore,
                                 bodLoad, bodConcentration, bodScore, ecoliLoad, ecoliConcentration, ecoliScore, magnitudeScore), typeList)

    addMessage('Saving calculations back to feature class')
    arcpy.da.ExtendTable(table, 'OID@', outArray, 'oid', False)

def CalculateTotalRisk(links, nodes):
    addMessage('Calculating total risk')
    arcpy.AddField_management(links, 'wq_risk_score', 'FLOAT')
    arcpy.AddField_management(nodes, 'wq_risk_score', 'FLOAT')

    arcpy.CalculateField_management(links, 'wq_risk_score', "!wq_susc_score! * !magnitude_score!", 'PYTHON')
    arcpy.CalculateField_management(nodes, 'wq_risk_score', "!wq_susc_score! * !magnitude_score!", 'PYTHON')


if __name__ == "__main__":
    try:
        outputDir = r'\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality'

        outputGdb = os.path.join(outputDir, 'WQRiskNetwork.gdb')
        linksBase = os.path.join(outputGdb, 'WQConveyance_base')
        nodesBase = os.path.join(outputGdb, 'WQNodes_base')
        inletCatchments = os.path.join(outputGdb, 'InletCatchmentsDetailed')
        flowDir = os.path.join(outputDir, 'flowdir_tau.tif')

        #arcpy.env.workspace = outputDir
        arcpy.env.overwriteOutput = True
        arcpy.env.snapRaster = flowDir
        arcpy.env.extent = arcpy.Describe(flowDir).Extent
        arcpy.CheckOutExtension("Spatial")

        runoffCoeffRasterPath = r'\\ws18325\Stormwater D\RiskModel\Surface\LandClassification.gdb\RunoffCoefficient'
        landuseRasterPath = r"\\ws18325\Stormwater D\RiskModel\Surface\LanduseZoningHybrid.tif"
        loadRasterGdb = r"\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality\LoadingEstimates.gdb"

        #Create the WQ base layers and accumulate land use area
        CreateWQNetwork(outputGdb, 'WQConveyance_base', 'WQNodes_base', inletCatchments, flowDir, landuseRasterPath, runoffCoeffRasterPath)
        rawq_susceptibility.AssignSusceptibility(linksBase, nodesBase)

        #Copy the base load network to use later for net load prior to running base load calculation
        linksNet = os.path.join(outputGdb, 'WQConveyance_net')
        nodesNet = os.path.join(outputGdb, 'WQNodes_net')

        arcpy.Copy_management(linksBase, linksNet)
        arcpy.Copy_management(nodesBase, nodesNet)

        #Calculate base loading for links and nodes
        CalculateLoading(linksBase)
        CalculateLoading(nodesBase)

        rawq_decumulate.CreateLoadRasters(linksNet, nodesNet, inletCatchments, flowDir, landuseRasterPath, loadRasterGdb)
        rawq_decumulate.CalculateBaseFlowAndVolume(linksNet, pd.annualRainfallDepth, 365)
        rawq_decumulate.CalculateBaseFlowAndVolume(nodesNet, pd.annualRainfallDepth, 365)

        rawq_decumulate.Decumulate(linksNet, nodesNet, inletCatchments)

        rawq_decumulate.CalculateScoresAndConcentrations(linksNet, 365)
        rawq_decumulate.CalculateScoresAndConcentrations(nodesNet, 365)

    except Exception as e:
        addMessage("Error: {0}".format(e.message))
    finally:
        addMessage("RAWQ script complete.")
